# 预训练器
from typing import List, Tuple
import itertools
import re
import json
import jsonlines
import psutil
import ujson
import numpy as np
import pandas as pd
from transformers import AutoTokenizer
from datasets import load_dataset
import os
from tqdm import tqdm
#
from nanods.nanods_tokenizer import NanodsTokenizer

class NanodsDataCurator(object):
    def __init__(self):
        self.name = 'nanods.nanods_data_curator.NanodsDataCurator'

    @staticmethod
    def load_tokenizer() -> NanodsTokenizer:
        # 定义BOS和EOS标记
        bos_token = "<s>"
        eos_token = "</s>"
        # 加载训练好的分词器路径
        tokenizer = AutoTokenizer.from_pretrained('./work/model/miniDeepSeek_tokenizer', use_fast=False)
        return tokenizer
    
    @staticmethod
    def preview_dataset(file_path:str, num_lines:int =5) -> None:
        """
        读取并展示数据集的前 num_lines 行
        """
        # 检查文件是否存在
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"{file_path} 文件不存在，请检查路径！")
        # 逐行读取并展示前 num_lines 行
        with jsonlines.open(file_path) as reader:
            for idx, obj in enumerate(reader):
                print(f"第 {idx + 1} 行数据: {obj}")
                if idx + 1 >= num_lines:
                    break

    @staticmethod
    def get_total_lines(file_path:str) -> int:
        """
        获取 JSONL 文件的总行数，不忽略错误，保证能够全面统计。
        """
        with open(file_path, 'rb') as f:  # 使用二进制模式避免编码问题
            return sum(1 for _ in f)
        
    @staticmethod
    def check_jsonl_format(file_path:str) -> Tuple[int, int, List]:
        """
        检查 JSONL 文件中的每一行是否是有效的 JSON 格式，带进度显示，并统计所有有问题的行。
        """
        total_lines = NanodsDataCurator.get_total_lines(file_path)  # 获取文件总行数
        valid_lines = 0
        invalid_lines = 0
        invalid_lns = [] # 非法行列表
        # 使用逐行读取，捕获 JSON 和编码错误
        with open(file_path, 'rb') as f:  # 使用二进制读取避免编码问题
            # 使用 tqdm 进度条显示检查进度
            for idx, line in tqdm(enumerate(f), total=total_lines, desc="Checking JSONL format"):
                try:
                    # 先尝试将每行数据解码为 UTF-8
                    decoded_line = line.decode('utf-8')
                    # 然后检查是否是有效的 JSON 格式
                    obj = jsonlines.Reader([decoded_line]).read()
                    valid_lines += 1
                except UnicodeDecodeError as e:
                    print(f"Encoding error at line {idx + 1}: {e}")
                    invalid_lns.append(idx+1)
                    invalid_lines += 1
                except jsonlines.InvalidLineError as e:
                    invalid_lns.append(idx+1)
                    print(f"Invalid JSON at line {idx + 1}: {e}")
                    invalid_lines += 1
        return valid_lines, invalid_lines, invalid_lns
    
    @staticmethod
    def remove_invalid_lines(file_path:str, output_path:str, invalid_lns:List) -> None:
        """
        读取文件，跳过指定的无效行，并将结果写入新文件
        """
        with open(file_path, 'rb') as infile, open(output_path, 'wb') as outfile:
            for idx, line in enumerate(infile):
                if idx + 1 not in invalid_lns:  # 跳过无效行
                    outfile.write(line)

    @staticmethod
    def process_seq_monkey(chunk_size=50000):
        """
        逐块读取 mobvoi_seq_monkey_general_open_corpus.jsonl 文件，
        对文本进行分词，并将分词结果保存为二进制文件，支持跳过无效行，并显示处理进度。
        """
        doc_ids = []
        chunk_idx = 0
        total_lines = 0
        # 定义BOS和EOS标记
        bos_token = "<s>"
        eos_token = "</s>"
        # 加载训练好的分词器路径
        tokenizer = AutoTokenizer.from_pretrained('./work/model/miniDeepSeek_tokenizer', use_fast=False)

        # 先计算总行数以便显示进度
        with open('./datasets/mobvoi_seq_monkey_general_open_corpus_cleaned.jsonl', 'r', encoding='utf-8') as f:
            total_lines = sum(1 for _ in f)

        # 打开jsonlines文件逐行读取
        with jsonlines.open('./datasets/mobvoi_seq_monkey_general_open_corpus_cleaned.jsonl') as reader:
            # 使用 tqdm 进度条显示进度
            with tqdm(total=total_lines, desc="Processing lines") as pbar:
                while True:
                    try:
                        # 使用 itertools.islice 按块读取文件，每次读取 chunk_size 行数据
                        chunk = list(itertools.islice(reader, chunk_size))
                    except jsonlines.InvalidLineError as e:
                        print(f"Skipping invalid chunk at chunk {chunk_idx}: {e}")
                        continue

                    if not chunk:  # 如果读取到文件末尾，则停止
                        break

                    # 遍历块中的每一行数据
                    # 逐行对数据进行编码（按token进行编码）
                    for idx, obj in enumerate(chunk):
                        try:
                            # 从每一行数据中提取'text'字段（即文本内容）
                            content = obj.get('text', '')
                            
                            # 跳过长度超过512的文本
                            if len(content) > 512:
                                continue

                            # 对文本进行分词，将其转为 token ids 序列，并加上BOS和EOS标记
                            text_id = tokenizer(f'{bos_token}{content}{eos_token}').data['input_ids']
                            
                            # 将分词结果添加到 doc_ids 列表中
                            doc_ids += text_id

                        except UnicodeDecodeError as e:
                            # 如果遇到编码错误，跳过该行，并打印错误信息
                            print(f"Skipping invalid line {chunk_idx * chunk_size + idx + 1}: {e}")
                            continue

                    # 每处理完一块数据，更新 chunk_idx 并打印进度信息
                    chunk_idx += 1
                    pbar.update(len(chunk))  # 更新进度条

                    # 如果累积的 token ids 超过 1,000,000 个，保存到文件中
                    if len(doc_ids) > 1000000:
                        arr = np.array(doc_ids, dtype=np.uint16)
                        with open(f'./datasets/clean_seq_monkey.bin', 'ab') as f:
                            f.write(arr.tobytes())
                        doc_ids = []

        # 如果处理完所有数据后 doc_ids 中还有未保存的内容，最后再保存一次
        if doc_ids:
            arr = np.array(doc_ids, dtype=np.uint16)
            with open(f'./datasets/clean_seq_monkey.bin', 'ab') as f:
                f.write(arr.tobytes())

    @staticmethod
    def pretrain_process():
        """
        函数的作用是调用 process_seq_monkey() 函数生成数据，
        然后整合所有生成的二进制文件，并将其合并保存为一个总的预训练数据文件。
        """
        # 调用 process_seq_monkey 函数处理数据
        NanodsDataCurator.process_seq_monkey()

        # 数据文件路径列表，目前只处理 clean_seq_monkey.bin 文件
        data_path_list = [
            './datasets/clean_seq_monkey.bin'
        ]
        
        data_lst = []
        
        # 读取生成的二进制文件
        for data_path in data_path_list:
            with open(data_path, 'rb') as f:
                # 将二进制文件中的内容加载到 numpy 数组中
                data = np.fromfile(f, dtype=np.uint16)
                data_lst.append(data)

        # 将所有读取到的数据合并为一个大数组
        arr = np.concatenate(data_lst)
        print(f"合并后的数据大小: {arr.shape}")

        # 将合并后的数据保存为最终的预训练数据文件
        with open('./datasets/pretrain_data.bin', 'wb') as f:
            f.write(arr.tobytes())

    @staticmethod
    def check_ds() -> None:
        # 读取CSV文件
        pretrain_df = pd.read_csv('./datasets/pretrain_data.csv')
        # 打印前5行数据（默认）
        print(pretrain_df.head())
        # 如果你想打印前n行，可以传递一个参数
        print(pretrain_df.head(n=10))  # 打印前10行